import Router from 'koa-router';
import { PassThrough } from 'stream';
import { generate } from '../utils/index.js';

const completion = new Router();

completion.post('/', async (ctx) => {
    ctx.set({
        'Connection': 'keep-alive',
        'Cache-Control': 'no-cache',
        'Content-Type': 'text/event-stream',
    });

    const stream = new PassThrough();
    ctx.body = stream;
    ctx.status = 200;

    const data = ctx.request.body;
    const prompt = `请补全下面的段落：\n${data.text}`;

    generate(prompt, [], (e) => {
        stream.write(JSON.stringify(e));
    }, () => {
        stream.end();
    });
});

export default completion;
